'''
https://leetcode.cn/problems/design-search-autocomplete-system
'''
import heapq
from typing import List


class Node:
    def __init__(self, sentence = None, times = 0):
        self.next = {}
        self.times = times
        self.sentence = sentence

    def __lt__(self, other):
        if self.times == other.times:
            return self.sentence > other.sentence
        return self.times < other.times

class AutocompleteSystem:

    def __init__(self, sentences: List[str], times: List[int]):
        self.root = Node()
        for word, time in zip(sentences, times):
            self.add_word(word, time)

        self.cur_word = ''
        self.cur = self.root

    def add_word(self, sentence, time):
        cur = self.root
        for c in sentence:
            if c not in cur.next:
                cur.next[c] = Node()
            cur = cur.next[c]
        cur.times += time
        cur.sentence = sentence

    def input(self, c: str) -> List[str]:
        if c == '#':
            self.cur.times += 1
            self.cur.sentence = self.cur_word
            self.cur_word = ''
            self.cur = self.root
            return []
        if c not in self.cur.next:
            self.cur.next[c] = Node()
        self.cur_word += c
        self.cur = self.cur.next[c]
        pq = []
        def dfs(root):
            if root.times:
                heapq.heappush(pq, root)
                if len(pq) > 3:
                    heapq.heappop(pq)
            for child in root.next.values():
                dfs(child)
        dfs(self.cur)
        res = [''] * (len(pq))
        for i in range(len(pq)-1, -1, -1):
            res[i] = heapq.heappop(pq).sentence
        return res

# Your AutocompleteSystem object will be instantiated and called as such:
sentences = ["i love you","island","iroman","i love leetcode"]
times = [5,3,2,2]
obj = AutocompleteSystem(sentences, times)
print(obj.input('i'))
print(obj.input(' '))
print(obj.input('a'))
print(obj.input('#'))
